import torch.nn as nn
# nn.functional工具包装载了网络层中那些只进行计算, 而没有参数的层
import torch.nn.functional as F
from learn_decoder import de_result

# 将线性层和softmax计算层一起实现, 因为二者的共同目标是生成最后的结构
# 因此把类的名字叫做Generator, 生成器类
class Generator(nn.Module):
    def __init__(self, d_model, vocab_size):
        """初始化函数的输入参数有两个, d_model代表词嵌入维度, vocab_size代表词表大小."""
        super(Generator, self).__init__()
        # 首先就是使用nn中的预定义线性层进行实例化, 得到一个对象self.project等待使用,
        # 这个线性层的参数有两个, 就是初始化函数传进来的两个参数：d_model, vocab_size
        self.project = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        """前向逻辑函数中输入是上一层的输出张量x"""
        # 在函数中, 首先使用上一步得到的self.project对x进行线性变化,
        # 然后使用F中已经实现的log_softmax进行的softmax处理.
        # 在这里之所以使用log_softmax是因为和我们这个pytorch版本的损失函数实现有关, 在其他版本中将修复.
        # log_softmax就是对softmax的结果又取了对数, 因为对数函数是单调递增函数,
        # 因此对最终我们取最大的概率值没有影响. 最后返回结果即可.
        return F.log_softmax(self.project(x), dim=-1)

# 例化参数:
# 词嵌入维度是512维
d_model = 512
# 词表大小是1000
vocab_size = 1000

# 输入参数:
# 输入x是上一层网络的输出, 我们使用来自解码器层的输出
x = de_result
print("de_result:", de_result, de_result.shape)
'''
de_result: tensor([[[ 0.3541, -0.1841,  0.6511,  ..., -0.3214, -1.2122,  0.3312],
         [-2.5232,  0.7296,  0.3281,  ..., -0.4447,  0.8314,  1.3124],
         [ 0.0892, -0.4234, -0.9003,  ..., -0.1831,  1.2330,  0.0580],
         [ 0.8192,  1.4067, -0.1367,  ..., -1.1207, -0.3351,  2.0668]],

        [[-0.0138, -0.7653,  0.1112,  ...,  2.0359,  0.5526,  0.2047],
         [ 0.5991,  0.2890, -0.1164,  ..., -0.2511, -0.0153,  0.0989],
         [ 1.6001, -0.2757,  0.2182,  ...,  0.0847, -0.9643, -1.6413],
         [ 0.4453,  0.1009,  2.3565,  ..., -0.1110,  0.0460,  0.2369]]],
       grad_fn=<AddBackward0>) torch.Size([2, 4, 512])'''

# 调用:
gen = Generator(d_model, vocab_size)
gen_result = gen(x)
print("gen_result:", gen_result, gen_result.shape)
'''
gen_result: tensor([[[-6.1963, -7.0450, -6.3052,  ..., -6.7417, -7.3108, -6.7984],
         [-6.3465, -6.5312, -8.2493,  ..., -6.6958, -7.5140, -6.5818],
         [-6.6637, -6.6967, -6.7903,  ..., -7.3226, -6.5956, -7.5328],
         [-8.6391, -7.7627, -7.6634,  ..., -7.4436, -7.4931, -7.2279]],

        [[-7.5593, -6.9726, -7.0190,  ..., -7.1281, -5.8868, -7.7423],
         [-6.6491, -7.3624, -9.0261,  ..., -7.4481, -6.6856, -6.4045],
         [-7.1506, -6.9171, -8.0529,  ..., -7.4694, -6.5181, -7.7354],
         [-7.8962, -6.8131, -7.7720,  ..., -8.1055, -7.4766, -5.4381]]],
       grad_fn=<LogSoftmaxBackward0>) torch.Size([2, 4, 1000])
'''